#!/usr/bin/env python3
"""
PyG 算子性能测试脚本

测试 /home/yzq/debug_RFdiffusion_pyg/ops/src 中的 PyG 算子性能
"""

import time
import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
import numpy as np
import os
import sys
import json
from typing import Dict, List, Tuple, Any
import psutil
import gc

# 添加路径以导入自定义算子
current_dir = os.path.dirname(os.path.abspath(__file__))
src_dir = os.path.join(current_dir, 'src')
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)
if src_dir not in sys.path:
    sys.path.insert(0, src_dir)

# 导入所有算子
try:
    from gspmm.gspmm import gspmm
    from gsddmm.gsddmm import gsddmm
    from edge_softmax.edge_softmax import edge_softmax
    from e_dot_v.e_dot_v import e_dot_v
    from copy_e_sum.copy_e_sum import copy_e_sum
    from pool.pooling import AvgPooling, MaxPooling
    from transform.remove_self_loops import remove_self_loops
    from graph.graph import Graph
except ImportError:
    # 备用导入方式
    sys.path.insert(0, os.path.join(current_dir, 'src'))
    from gspmm.gspmm import gspmm
    from gsddmm.gsddmm import gsddmm  
    from edge_softmax.edge_softmax import edge_softmax
    from e_dot_v.e_dot_v import e_dot_v
    from copy_e_sum.copy_e_sum import copy_e_sum
    from pool.pooling import AvgPooling, MaxPooling
    from transform.remove_self_loops import remove_self_loops
    from graph.graph import Graph


class PerformanceTester:
    """PyG 算子性能测试器"""
    
    def __init__(self, device='cuda', warmup_iterations=5, test_iterations=10):
        """
        初始化性能测试器
        
        Args:
            device: 设备类型 ('cuda', 'cpu', 'npu')
            warmup_iterations: 热身迭代次数
            test_iterations: 测试迭代次数
        """
        self.device = torch.device(device)
        self.warmup_iterations = warmup_iterations
        self.test_iterations = test_iterations
        self.results = {}
        
        # 池化算子
        self.avg_pool = AvgPooling()
        self.max_pool = MaxPooling()
        
        print(f"初始化性能测试器 - 设备: {self.device}")
        if device == 'npu':
            print(f"NPU 设备数量: {torch.npu.device_count()}")
        elif device == 'cuda':
            print(f"CUDA 设备数量: {torch.cuda.device_count()}")
    
    def create_test_graph(self, num_nodes: int, num_edges: int, 
                         node_features: int = 64, edge_features: int = 32) -> Tuple[Graph, torch.Tensor, torch.Tensor]:
        """
        创建测试图和特征
        
        Args:
            num_nodes: 节点数量
            num_edges: 边数量
            node_features: 节点特征维度
            edge_features: 边特征维度
            
        Returns:
            图对象，节点特征，边特征
        """
        # 随机生成边索引
        src = torch.randint(0, num_nodes, (num_edges,), device=self.device)
        dst = torch.randint(0, num_nodes, (num_edges,), device=self.device)
        
        # 创建图
        graph = Graph((src, dst), num_nodes=num_nodes)
        graph = graph.to(self.device)
        
        # 创建节点特征
        node_feat = torch.randn(num_nodes, node_features, device=self.device)
        
        # 创建边特征
        edge_feat = torch.randn(num_edges, edge_features, device=self.device)
        
        return graph, node_feat, edge_feat
    
    def measure_memory(self) -> float:
        """测量当前内存使用量 (MB)"""
        if self.device.type == 'cuda':
            return torch.cuda.memory_allocated() / 1024 / 1024
        elif self.device.type == 'npu':
            return torch.npu.memory_allocated() / 1024 / 1024
        else:
            return psutil.Process().memory_info().rss / 1024 / 1024
    
    def time_function(self, func, *args, **kwargs) -> Tuple[float, Any]:
        """
        测量函数执行时间
        
        Args:
            func: 要测试的函数
            *args, **kwargs: 函数参数
            
        Returns:
            (平均执行时间(ms), 函数结果)
        """
        # 热身
        for _ in range(self.warmup_iterations):
            try:
                result = func(*args, **kwargs)
                if self.device.type in ['cuda', 'npu']:
                    torch.cuda.synchronize() if self.device.type == 'cuda' else torch.npu.synchronize()
            except Exception as e:
                print(f"热身阶段出错: {e}")
                return float('inf'), None
        
        # 测试
        times = []
        result = None
        
        for _ in range(self.test_iterations):
            if self.device.type in ['cuda', 'npu']:
                torch.cuda.synchronize() if self.device.type == 'cuda' else torch.npu.synchronize()
            
            start_time = time.perf_counter()
            
            try:
                result = func(*args, **kwargs)
                
                if self.device.type in ['cuda', 'npu']:
                    torch.cuda.synchronize() if self.device.type == 'cuda' else torch.npu.synchronize()
                
                end_time = time.perf_counter()
                times.append((end_time - start_time) * 1000)  # 转换为毫秒
                
            except Exception as e:
                print(f"测试执行出错: {e}")
                return float('inf'), None
        
        avg_time = np.mean(times)
        return avg_time, result
    
    def test_gspmm(self, graph: Graph, node_feat: torch.Tensor, edge_feat: torch.Tensor) -> Dict[str, float]:
        """测试 GSPMM 算子性能"""
        print("测试 GSPMM 算子...")
        results = {}
        
        # 测试不同的操作和聚合方式
        operations = ['mul', 'add', 'sub']
        reduce_ops = ['sum', 'max', 'min', 'mean']
        
        for op in operations:
            for reduce_op in reduce_ops:
                test_name = f"gspmm_{op}_{reduce_op}"
                print(f"  测试 {test_name}")
                
                try:
                    avg_time, _ = self.time_function(
                        gspmm, graph, op, reduce_op, edge_feat, node_feat
                    )
                    results[test_name] = avg_time
                except Exception as e:
                    print(f"    错误: {e}")
                    results[test_name] = float('inf')
        
        return results
    
    def test_gsddmm(self, graph: Graph, node_feat: torch.Tensor) -> Dict[str, float]:
        """测试 GSDDMM 算子性能"""
        print("测试 GSDDMM 算子...")
        results = {}
        
        # 测试不同的操作
        operations = ['mul', 'add', 'sub', 'div', 'dot']
        targets = [('u', 'v'), ('u', 'u'), ('v', 'v')]
        
        for op in operations:
            for lhs_target, rhs_target in targets:
                test_name = f"gsddmm_{op}_{lhs_target}_{rhs_target}"
                print(f"  测试 {test_name}")
                
                try:
                    avg_time, _ = self.time_function(
                        gsddmm, graph, op, node_feat, node_feat, lhs_target, rhs_target
                    )
                    results[test_name] = avg_time
                except Exception as e:
                    print(f"    错误: {e}")
                    results[test_name] = float('inf')
        
        return results
    
    def test_edge_softmax(self, graph: Graph, edge_feat: torch.Tensor) -> Dict[str, float]:
        """测试边 Softmax 算子性能"""
        print("测试 Edge Softmax 算子...")
        results = {}
        
        # 测试一维和多维边特征
        test_cases = [
            ("1d", edge_feat[:, 0]),  # 一维
            ("2d", edge_feat),        # 二维
        ]
        
        for name, feat in test_cases:
            test_name = f"edge_softmax_{name}"
            print(f"  测试 {test_name}")
            
            try:
                avg_time, _ = self.time_function(edge_softmax, graph, feat)
                results[test_name] = avg_time
            except Exception as e:
                print(f"    错误: {e}")
                results[test_name] = float('inf')
        
        return results
    
    def test_e_dot_v(self, graph: Graph, node_feat: torch.Tensor, edge_feat: torch.Tensor) -> Dict[str, float]:
        """测试 E dot V 算子性能"""
        print("测试 E dot V 算子...")
        
        try:
            avg_time, _ = self.time_function(e_dot_v, graph, edge_feat, node_feat)
            return {"e_dot_v": avg_time}
        except Exception as e:
            print(f"  错误: {e}")
            return {"e_dot_v": float('inf')}
    
    def test_copy_e_sum(self, graph: Graph, edge_feat: torch.Tensor) -> Dict[str, float]:
        """测试 Copy E Sum 算子性能"""
        print("测试 Copy E Sum 算子...")
        results = {}
        
        # 测试一维和多维边特征
        test_cases = [
            ("1d", edge_feat[:, 0]),  # 一维
            ("2d", edge_feat),        # 二维
        ]
        
        for name, feat in test_cases:
            test_name = f"copy_e_sum_{name}"
            print(f"  测试 {test_name}")
            
            try:
                avg_time, _ = self.time_function(copy_e_sum, graph, feat)
                results[test_name] = avg_time
            except Exception as e:
                print(f"    错误: {e}")
                results[test_name] = float('inf')
        
        return results
    
    def test_pooling(self, graph: Graph, node_feat: torch.Tensor) -> Dict[str, float]:
        """测试池化算子性能"""
        print("测试池化算子...")
        results = {}
        
        # 平均池化
        try:
            avg_time, _ = self.time_function(self.avg_pool, graph, node_feat)
            results["avg_pooling"] = avg_time
        except Exception as e:
            print(f"  平均池化错误: {e}")
            results["avg_pooling"] = float('inf')
        
        # 最大池化
        try:
            avg_time, _ = self.time_function(self.max_pool, graph, node_feat)
            results["max_pooling"] = avg_time
        except Exception as e:
            print(f"  最大池化错误: {e}")
            results["max_pooling"] = float('inf')
        
        return results
    
    def test_remove_self_loops(self, graph: Graph) -> Dict[str, float]:
        """测试移除自环算子性能"""
        print("测试移除自环算子...")
        
        # 为图添加一些自环边
        num_self_loops = min(100, graph.num_nodes())
        self_loop_nodes = torch.randint(0, graph.num_nodes(), (num_self_loops,), device=self.device)
        
        # 添加自环
        src_with_loops = torch.cat([graph.edge_index[0], self_loop_nodes])
        dst_with_loops = torch.cat([graph.edge_index[1], self_loop_nodes])
        
        graph_with_loops = Graph((src_with_loops, dst_with_loops), num_nodes=graph.num_nodes())
        graph_with_loops = graph_with_loops.to(self.device)
        
        try:
            avg_time, _ = self.time_function(remove_self_loops, graph_with_loops)
            return {"remove_self_loops": avg_time}
        except Exception as e:
            print(f"  错误: {e}")
            return {"remove_self_loops": float('inf')}
    
    def run_comprehensive_test(self, test_configs: List[Dict[str, int]]) -> Dict[str, Any]:
        """
        运行全面的性能测试
        
        Args:
            test_configs: 测试配置列表，每个配置包含图的大小参数
            
        Returns:
            测试结果字典
        """
        print("开始综合性能测试...")
        print("=" * 60)
        
        all_results = {}
        
        for i, config in enumerate(test_configs):
            config_name = f"config_{i+1}_n{config['num_nodes']}_e{config['num_edges']}"
            print(f"\n配置 {i+1}/{len(test_configs)}: {config}")
            print("-" * 40)
            
            # 清理内存
            gc.collect()
            if self.device.type == 'cuda':
                torch.cuda.empty_cache()
            elif self.device.type == 'npu':
                torch.npu.empty_cache()
            
            # 创建测试图
            graph, node_feat, edge_feat = self.create_test_graph(**config)
            
            # 记录内存使用
            memory_before = self.measure_memory()
            
            config_results = {
                'config': config,
                'memory_before_mb': memory_before,
                'tests': {}
            }
            
            # 运行各个算子测试
            config_results['tests'].update(self.test_gspmm(graph, node_feat, edge_feat))
            config_results['tests'].update(self.test_gsddmm(graph, node_feat))
            config_results['tests'].update(self.test_edge_softmax(graph, edge_feat))
            config_results['tests'].update(self.test_e_dot_v(graph, node_feat, edge_feat))
            config_results['tests'].update(self.test_copy_e_sum(graph, edge_feat))
            config_results['tests'].update(self.test_pooling(graph, node_feat))
            config_results['tests'].update(self.test_remove_self_loops(graph))
            
            # 记录测试后内存使用
            memory_after = self.measure_memory()
            config_results['memory_after_mb'] = memory_after
            config_results['memory_diff_mb'] = memory_after - memory_before
            
            all_results[config_name] = config_results
            
            print(f"内存使用: {memory_before:.2f}MB -> {memory_after:.2f}MB (差值: {memory_after - memory_before:.2f}MB)")
        
        return all_results
    
    def generate_report(self, results: Dict[str, Any], output_file: str = "performance_report.json"):
        """生成性能测试报告"""
        print("\n" + "=" * 60)
        print("性能测试报告")
        print("=" * 60)
        
        # 汇总结果
        summary = {
            'device': str(self.device),
            'warmup_iterations': self.warmup_iterations,
            'test_iterations': self.test_iterations,
            'results': results,
            'summary_statistics': {}
        }
        
        # 计算统计信息
        all_times = []
        operator_times = {}
        
        for config_name, config_results in results.items():
            print(f"\n配置: {config_name}")
            print(f"  图大小: {config_results['config']['num_nodes']} 节点, {config_results['config']['num_edges']} 边")
            print(f"  内存使用: {config_results['memory_diff_mb']:.2f}MB")
            print("  算子性能 (毫秒):")
            
            for test_name, time_ms in config_results['tests'].items():
                if time_ms != float('inf'):
                    all_times.append(time_ms)
                    if test_name not in operator_times:
                        operator_times[test_name] = []
                    operator_times[test_name].append(time_ms)
                    print(f"    {test_name}: {time_ms:.3f}ms")
                else:
                    print(f"    {test_name}: 失败")
        
        # 汇总统计
        if all_times:
            summary['summary_statistics'] = {
                'total_tests': len(all_times),
                'average_time_ms': np.mean(all_times),
                'median_time_ms': np.median(all_times),
                'min_time_ms': np.min(all_times),
                'max_time_ms': np.max(all_times),
                'std_time_ms': np.std(all_times)
            }
            
            print(f"\n总体统计:")
            print(f"  总测试数: {len(all_times)}")
            print(f"  平均耗时: {np.mean(all_times):.3f}ms")
            print(f"  中位数耗时: {np.median(all_times):.3f}ms")
            print(f"  最短耗时: {np.min(all_times):.3f}ms")
            print(f"  最长耗时: {np.max(all_times):.3f}ms")
            print(f"  标准差: {np.std(all_times):.3f}ms")
        
        # 算子性能排名
        print(f"\n算子平均性能排名:")
        operator_avg_times = {op: np.mean(times) for op, times in operator_times.items() if times}
        sorted_operators = sorted(operator_avg_times.items(), key=lambda x: x[1])
        
        for i, (op, avg_time) in enumerate(sorted_operators[:10]):  # 显示前10个
            print(f"  {i+1}. {op}: {avg_time:.3f}ms")
        
        # 保存结果到文件
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(summary, f, indent=2, ensure_ascii=False)
        
        print(f"\n详细结果已保存到: {output_file}")


def main():
    """主函数"""
    print("PyG 算子性能测试脚本")
    print("=" * 60)
    
    # 检测可用设备
    available_devices = ['cpu']
    if torch.cuda.is_available():
        available_devices.append('cuda')
    if hasattr(torch, 'npu') and torch.npu.is_available():
        available_devices.append('npu')
    
    print(f"可用设备: {available_devices}")
    
    # 选择设备（优先级：npu > cuda > cpu）
    if 'npu' in available_devices:
        device = 'npu'
    elif 'cuda' in available_devices:
        device = 'cuda'
    else:
        device = 'cpu'
    
    print(f"使用设备: {device}")
    
    # 测试配置 - 不同规模的图
    test_configs = [
        # 小规模图
        {'num_nodes': 100, 'num_edges': 500, 'node_features': 32, 'edge_features': 16},
        {'num_nodes': 500, 'num_edges': 2000, 'node_features': 64, 'edge_features': 32},
        
        # 中等规模图
        {'num_nodes': 1000, 'num_edges': 5000, 'node_features': 64, 'edge_features': 32},
        {'num_nodes': 2000, 'num_edges': 10000, 'node_features': 128, 'edge_features': 64},
        
        # 大规模图
        {'num_nodes': 5000, 'num_edges': 25000, 'node_features': 128, 'edge_features': 64},
        {'num_nodes': 10000, 'num_edges': 50000, 'node_features': 256, 'edge_features': 128},
    ]
    
    # 创建性能测试器
    tester = PerformanceTester(device=device, warmup_iterations=3, test_iterations=5)
    
    # 运行测试
    results = tester.run_comprehensive_test(test_configs)
    
    # 生成报告
    output_file = f"pyg_ops_performance_report_{device}.json"
    tester.generate_report(results, output_file)
    
    print("\n测试完成！")


if __name__ == "__main__":
    main() 